import pandas as pd
from openai import OpenAI
import time
import re
import numpy as np
from concurrent.futures import ThreadPoolExecutor, as_completed
import json

# 初始化 DeepSeek 客户端
client = OpenAI(
    api_key="sk-b58015698aa548c0ba62f8d30fbfd06f",
    base_url="https://api.deepseek.com/v1"
)

# 预定义列名翻译映射
COLUMN_TRANSLATIONS = {
    'brand': '品牌',
    'model': '型号',
    'top_speed_kmh': '最高时速(公里/小时)',
    'battery_capacity_kWh': '电池容量(kWh)',
    'torque_nm': '扭矩(牛·米)',
    'efficiency_wh_per_km': '能耗(瓦时/公里)',
    'range_km': '续航里程(公里)',
    'acceleration_0_100_s': '0-100公里/小时加速时间(秒)',
    'fast_charging_power_kw_dc': '快充功率(kW)',
    'towing_capacity_kg': '牵引能力(kg)',
    'cargo_volume_l': '货舱容积(L)',
    'seats': '座位数',
    'drivetrain': '传动系统',
    'segment': '车辆类别',
    'length_m': '长度(m)',
    'width_m': '宽度(m)',
    'height_m': '高度(m)',
    'car_body_type': '车身类型'
}

# 汽车专业术语缓存
car_term_cache = {
    # 驱动系统
    "AWD": "四驱",
    "RWD": "后驱",
    "FWD": "前驱",

    # 车身类型
    "Hatchback": "掀背车",
    "Station/Estate": "旅行车",
    "Liftback Sedan": "掀背轿车",
    "SUV": "SUV",
    "Sedan": "轿车",
    "Small Passenger Van": "小型客车",
    "Crossover": "跨界车",
    "Coupe": "双门轿跑",
    "Convertible": "敞篷车",
    "Minivan": "小型货车",
    "Pickup": "皮卡",
    "MPV": "MPV",

    # 单位
    "kW": "千瓦",
    "kWh": "千瓦时",
    "N·m": "牛·米",
    "Wh/km": "瓦时/公里",
    "km/h": "公里/小时",
    "s": "秒",
    "kg": "千克",
    "L": "升",
    "m": "米",

    # 技术术语
    "EV": "电动汽车",
    "PHEV": "插电式混合动力",
    "BEV": "纯电动汽车",
    "HEV": "混合动力",
    "FCV": "燃料电池汽车",
    "DCFC": "直流快充",
    "Regen": "能量回收",
    "Frunk": "前行李箱",
    "Autopilot": "自动驾驶系统",
    "ADAS": "高级驾驶辅助系统"
}

# 汽车品牌名称翻译映射（根据您提供的数据）
brand_translation = {
    "Alpine": "Alpine",
    "Audi": "奥迪",
    "BMW": "宝马",
    "BYD": "比亚迪",
    "CUPRA": "CUPRA",
    "Cadillac": "凯迪拉克",
    "Citroen": "雪铁龙",
    "DS": "DS汽车",
    "Fiat": "菲亚特",
    "Ford": "福特",
    "Genesis": "捷尼赛思",
    "Hongqi": "红旗",
    "Hyundai": "现代",
    "Jaguar": "捷豹",
    "Jeep": "吉普",
    "KGM": "KGM",
    "Kia": "起亚",
    "Leapmotor": "零跑",
    "Lexus": "雷克萨斯",
    "Lotus": "路特斯",
    "Lynk&Co": "领克",
    "MG": "名爵",
    "Maserati": "玛莎拉蒂",
    "Maxus": "上汽大通",
    "Mazda": "马自达",
    "Mercedes-Benz": "梅赛德斯-奔驰",
    "Mini": "迷你",
    "NIO": "蔚来",
    "Nissan": "日产",
    "Opel": "欧宝",
    "Peugeot": "标致",
    "Polestar": "极星",
    "Porsche": "保时捷",
    "Renault": "雷诺",
    "Skoda": "斯柯达",
    "Smart": "smart",
    "Subaru": "斯巴鲁",
    "Tesla": "特斯拉",
    "Toyota": "丰田",
    "VinFast": "VinFast",
    "Volkswagen": "大众",
    "Volvo": "沃尔沃",
    "Voyah": "岚图",
    "XPENG": "小鹏",
    "Zeekr": "极氪"
}

# 车辆类别翻译映射（根据您提供的数据）
segment_translation = {
    "B - Compact": "B级-紧凑型",
    "JE - Executive": "JE级-行政型",
    "JC - Medium": "JC级-中型",
    "JD - Large": "JD级-大型",
    "D - Large": "D级-大型",
    "E - Executive": "E级-行政型",
    "F - Luxury": "F级-豪华型",
    "JB - Compact": "JB级-紧凑型",
    "JF - Luxury": "JF级-豪华型",
    "N - Passenger Van": "N级-乘用货车"
}


def translate_with_r1(text, target_lang="中文", max_retries=3):
    """使用DeepSeek-R1模型翻译文本，带缓存和重试机制"""
    if pd.isna(text) or not str(text).strip():
        return text

    # 检查缓存
    text_str = str(text).strip()
    if text_str in car_term_cache:
        return car_term_cache[text_str]

    # 特殊处理：品牌名称
    if text_str in brand_translation:
        return brand_translation[text_str]

    # 特殊处理：车辆类别
    if text_str in segment_translation:
        return segment_translation[text_str]

    # 特殊处理：数字和单位组合 (e.g., "150kW", "75kWh")
    if re.match(r'^\d+\.?\d*\s*[a-zA-Z·/]+$', text_str):
        parts = re.split(r'(\d+\.?\d*)', text_str)
        if len(parts) >= 3:
            number = parts[1]
            unit = parts[2].strip()
            if unit in car_term_cache:
                return f"{number}{car_term_cache[unit]}"

    # 特殊处理：驱动系统缩写
    if text_str in ["AWD", "RWD", "FWD", "4WD"]:
        return car_term_cache.get(text_str, text_str)

    # 构造更明确的提示词
    prompt = (
        "你是一名汽车工程专业翻译，请准确翻译以下汽车相关术语：\n"
        f"### 原文:\n{text_str}\n\n"
        "### 翻译要求:\n"
        "1. 专业术语保持行业标准译法\n"
        "2. 品牌名称使用官方中文名\n"
        "3. 技术参数保持原格式\n"
        "4. 仅返回翻译结果，不加额外说明\n"
        "5. 如果不需要翻译或无法翻译，请直接返回原文\n"
        "### 中文翻译:"
    )

    retries = 0
    while retries < max_retries:
        try:
            response = client.chat.completions.create(
                model="deepseek-reasoner",
                messages=[
                    {"role": "system", "content": "你是一名专业的汽车行业翻译专家，精通中英互译"},
                    {"role": "user", "content": prompt}
                ],
                temperature=0.1,
                max_tokens=100,
                stream=False
            )

            translated = response.choices[0].message.content.strip()

            # 清理API返回的额外内容
            translated = re.split(r'[:：]', translated)[-1].strip()
            translated = translated.strip('"\'')

            # 如果翻译结果为空或与原文相同，返回原文
            if not translated or translated == text_str:
                return text_str

            # 缓存结果
            car_term_cache[text_str] = translated
            return translated
        except Exception as e:
            print(f"翻译出错 (尝试 {retries + 1}/{max_retries}): {e} | 原文: {text_str}")
            retries += 1
            time.sleep(1.5)  # 出错时等待

    # 重试失败后返回原文
    return text_str


def batch_translate(values, batch_size=30):
    """批量翻译函数，显著提升效率"""
    # 首先检查缓存
    cached_values = [car_term_cache.get(str(v), None) for v in values]
    to_translate = []
    indices = []

    # 找出需要翻译的值
    for i, (value, cached) in enumerate(zip(values, cached_values)):
        if cached is None:
            to_translate.append(str(value))
            indices.append(i)

    print(f"需要翻译 {len(to_translate)}/{len(values)} 个值")

    # 批量翻译
    translated_results = []
    for i in range(0, len(to_translate), batch_size):
        batch = to_translate[i:i + batch_size]
        with ThreadPoolExecutor(max_workers=8) as executor:
            futures = [executor.submit(translate_with_r1, value) for value in batch]
            for future in as_completed(futures):
                translated_results.append(future.result())
        print(f"已翻译: {min(i + batch_size, len(to_translate))}/{len(to_translate)}")
        time.sleep(0.8)  # 避免API速率限制

    # 更新缓存
    for idx, value, trans in zip(indices, to_translate, translated_results):
        car_term_cache[value] = trans
        cached_values[idx] = trans

    return cached_values


def translate_column(column):
    """高效翻译列数据"""
    print(f"翻译列: '{column.name}'，值数量: {len(column)}")

    # 处理唯一值，减少API调用
    unique_values = column.unique()
    unique_translated = batch_translate(unique_values)

    # 创建翻译映射
    translation_map = dict(zip(unique_values, unique_translated))

    # 应用映射
    return column.map(translation_map)


def main():
    # 读取数据
    df = pd.read_csv('dataset/electric_vehicles_spec_2025_clean.csv')

    print("=" * 50)
    print("开始翻译列名...")
    print("=" * 50)

    # 使用预定义的列名翻译
    new_columns = []
    for col in df.columns:
        # 使用预定义翻译，如果没有则调用API
        if col in COLUMN_TRANSLATIONS:
            translated = COLUMN_TRANSLATIONS[col]
        else:
            translated = translate_with_r1(col)
            COLUMN_TRANSLATIONS[col] = translated  # 缓存新翻译

        new_columns.append(translated)
        print(f"列名翻译: '{col}' -> '{translated}'")

    df.columns = new_columns

    print("\n" + "=" * 50)
    print("开始翻译数据内容...")
    print("=" * 50)

    # 仅翻译文本列
    text_columns = df.select_dtypes(include=['object']).columns

    # 特别处理品牌列 - 使用预定义的品牌翻译
    if '品牌' in df.columns:
        print("处理品牌列...")
        df['品牌'] = df['品牌'].map(lambda x: brand_translation.get(str(x).strip(), x))

    # 特别处理车辆类别列 - 使用预定义的segment翻译
    if '车辆类别' in df.columns:
        print("处理车辆类别列...")
        df['车辆类别'] = df['车辆类别'].map(lambda x: segment_translation.get(str(x).strip(), x))

    # 特别处理驱动系统列
    if '传动系统' in df.columns:
        print("处理传动系统列...")
        df['传动系统'] = df['传动系统'].map(lambda x: car_term_cache.get(str(x).strip(), x))

    # 特别处理车身类型列
    if '车身类型' in df.columns:
        print("处理车身类型列...")
        df['车身类型'] = df['车身类型'].map(lambda x: car_term_cache.get(str(x).strip(), x))

    # 跳过型号(model)列的翻译
    if '型号' in df.columns:
        print("跳过型号列的翻译，保留原文...")

    # 翻译其他文本列
    for col in text_columns:
        # 跳过已处理的列和型号列
        if col in ['品牌', '车辆类别', '传动系统', '车身类型', '型号']:
            continue

        print(f"翻译列: {col}")
        df[col] = translate_column(df[col])

    # 保存翻译结果
    output_path = 'dataset/electric_vehicles_spec_2025_zh.csv'
    df.to_csv(output_path, index=False, encoding='utf-8-sig')

    print("\n" + "=" * 50)
    print(f"翻译完成! 结果已保存至: {output_path}")
    print(f"总缓存条目: {len(car_term_cache)}")
    print("=" * 50)

    # 保存列名翻译映射供将来使用
    with open('column_translations.json', 'w', encoding='utf-8') as f:
        json.dump(COLUMN_TRANSLATIONS, f, ensure_ascii=False, indent=2)

    # 保存品牌翻译映射
    with open('brand_translations.json', 'w', encoding='utf-8') as f:
        json.dump(brand_translation, f, ensure_ascii=False, indent=2)

    # 保存车辆类别翻译
    with open('segment_translations.json', 'w', encoding='utf-8') as f:
        json.dump(segment_translation, f, ensure_ascii=False, indent=2)


if __name__ == "__main__":
    main()